import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib as mpl
from patsy import dmatrix
from sklearn.ensemble import RandomForestClassifier
from sklearn import linear_model



def get_ipo_data():
    """
    日期，发行者，股票代码，主承销商/联席主承销商， 发行价，开盘价， 首日收盘价
    首日价格变化比例: （收盘 - 发行价） / 发行价
    开盘价-发行价
    收盘价-发行价
    评级
    成交
    :return:
    """
    return pd.read_excel('./data/ipo_data.xlsx')


def get_sp500_data():
    """
    标普500指数
    日期，开盘，最高，最低，收盘，成交量，调整收盘价
    日变化：收盘-开盘
    隔夜变化：开盘 - 前收盘
    :return:
    """
    df = pd.read_excel('./data/sp500.xlsx')
    df.sort_values('Date', inplace=True)
    df.reset_index(drop=True, inplace=True)
    return df


def get_sp500_pct(df_spy, date):
    info = df_spy[df_spy['Date']==date]
    info_dict = {'Pct': info['Pct'], 'Pct_7': info['Pct_7']}
    return pd.Series(info_dict)


def get_data():
    """
    获取全部数据
    :return:
    """
    # 1、获取数据
    df_ipo = get_ipo_data()
    df_spy = get_sp500_data()
    # spy 7日收盘变化率
    df_spy['Pct_7'] = df_spy['Close'].pct_change(periods=7)
    df_spy['Pct'] = df_spy['Close'].pct_change()

    # 2、准备数据
    df_ipo = prepare(df_ipo)

    # 3、数据合并
    # 检查 sp500没数据，但是ipo有数据的情况 ipo有周末发行的？
    def get_pct(date):
        """
        获取对应数据
        :param date:
        :return:
        """
        row_dict = dict()

        date_info = df_spy[df_spy['Date'] == date]
        if not date_info.empty:
            row_dict['Pct'] = date_info['Pct'].iloc[0]
            row_dict['Pct_7'] = date_info['Pct_7'].iloc[0]
        return pd.Series(row_dict)

    def method_1():
        """
        方法一
        :return:
        """
        pct_info = df_ipo['Date'].apply(get_pct)
        df_new = pd.concat([df_ipo, pct_info], axis=1)

        df_new = df_new.sort_values('Date')

        df_new['Pct'] = df_new['Pct'].shift(1)
        df_new['Pct_7'] = df_new['Pct_7'].shift(1)

        pct, pct_7 = np.nan, np.nan
        date = None
        for index, row in df_new.iterrows():
            if date != row['Date']:
                pct = row['Pct']
                pct_7 = row['Pct_7']
            else:
                df_new.loc[index, 'Pct'] = pct
                df_new.loc[index, 'Pct_7'] = pct_7

            date = row['Date']

        # print(df_new[['Date', 'Pct', 'Pct_7']])
        return df_new

    def method_2(df_ipo):
        """
        方法二
        :return:
        """
        pct = [df_spy[df_spy['Date'] == x]['Pct'].iloc[0] if not df_spy[df_spy['Date'] == x].empty else np.nan
               for x in df_ipo['Date']]
        pct_7 = [np.nan if df_spy[df_spy['Date'] == x].empty else df_spy[df_spy['Date'] == x]['Pct_7'].iloc[0]
                 for x in df_ipo['Date']]

        df_ipo['Pct'] = pd.DataFrame(pct, index=df_ipo.index)
        df_ipo['Pct_7'] = pd.DataFrame(pct_7, index=df_ipo.index)

        df_ipo = df_ipo.sort_values('Date')

        df_ipo['Pct'] = df_ipo['Pct'].shift(1)
        df_ipo['Pct_7'] = df_ipo['Pct_7'].shift(1)

        pct, pct_7 = np.nan, np.nan
        date = None
        for index, row in df_ipo.iterrows():
            if date != row['Date']:
                pct = row['Pct']
                pct_7 = row['Pct_7']
            else:
                df_ipo.loc[index, 'Pct'] = pct
                df_ipo.loc[index, 'Pct_7'] = pct_7

            date = row['Date']

        return df_ipo

    return method_1()


def prepare(df):
    """
    数据准备
    :param df:
    :return:
    """
    print(df.head().T)
    # print(df.info())

    # 1、去掉异常值
    # df['Ratings'] = df['Ratings'].apply(lambda x: 0 if x == 'N/C' else int(x))
    df.replace('N/C', 0, inplace=True)

    # 2、拆分取主承销商
    df['Lead Mgr'] = df['Managers'].map(lambda x: x.split('/')[0])
    df['Lead Mgr'] = df['Lead Mgr'].map(lambda x: x.strip())

    # 3、投行(主承销商)处理
    bank_dict = {
        'SunTrust': 'SunTrust Robinson',
        'Wedbush': 'Wedbush Morgan',
        'Blair': 'William Blair',
        'Max': 'Maxim Group',
        'CRT': 'CRT Capital',
        'HCF': 'HCFP Brenner ',
        'Cohen': 'Cohen & Co.',
        'Cowen': 'Cowen & Co.',
        'Leerink': 'Leerink Partners',
        'Lynch\\xca': 'Merrill Lynch',
        'Merril ': 'Merrill Lynch',
        'BofA$': 'BofA Merrill Lynch',
        'SANDLER': "Sandler O'neil + Partners",
        'Sandler': "Sandler O'neil + Partners",
        'Renshaw': 'Rodman & Renshaw',
        'Baird': 'RW Baird',
        'Cantor': 'Cantor Fitzgerald',
        'Goldman': 'Goldman Sachs',
        'Bear': 'Bear Stearns',
        'BoA': 'BofA Merrill Lynch',
        'Broadband': 'Broadband Capital',
        'Davidson': 'DA Davidson',
        'Feltl': 'Feltl & Co.',
        'China': 'China International',
        'Cit': 'Citigroup',
        'Ferris': 'Ferris Baker Watts',
        'Friedman|Freidman|FBR': 'Friedman Billings Ramsey',
        '^I-': 'I-Bankers',
        'Gunn': 'Gunn Allen',
        'Jeffer': 'Jefferies',
        'Oppen': 'Oppenheimer',
        'JMP': 'JMP Securities',
        'Rice': 'Johnson Rice',
        'Ladenburg': 'Ladenburg Thalmann',
        'Piper': 'Piper Jaffray',
        'Pali': 'Pali Capital',
        'Paulson': 'Paulson Investment Co.',
        'Roth': 'Roth Capital',
        'Stifel': 'Stifel Nicolaus',
        'Hambrecht': 'WRHambrecht+Co.',
        'Edwards': 'AGEdwards',
        'Barclay': 'Barclays',
        'Aegis': 'Aegis Capital',
        'Deutsche': 'Deutsche Bank',
        'Suisse': 'CSFB',
        '^Early': 'EarlyBird Capital',
        'CS.?F': 'CSFB',
        'Keefe': 'Keefe, Bruyette & Woods',
        'Stan': 'Morgan Stanley',
        'P. Morg': 'JP Morgan',
        'PM': 'JP Morgan',
        'Banc of': 'Banc of America',
        'Lych': 'BofA Merrill Lynch',
        'Merrill$': 'Merrill Lynch',
        'Lymch': 'Merrill Lynch',
        'A Merril Lynch': 'BofA Merrill Lynch'
    }

    for k, v in bank_dict.items():
        df.loc[df['Lead Mgr'].str.contains(k), 'Lead Mgr'] = v

    df.loc[325, 'Lead Mgr'] = 'Maximum Captial'

    # 4、拆分取所有承销商数量
    df['Total Underwriters'] = df['Managers'].map(lambda x: len(x.split('/')))

    # 5、增加日期特征
    df['Week'] = df['Date'].apply(lambda x: x.date().weekday() + 1)
    df['Month'] = df['Date'].apply(lambda x: x.date().month)

    # 6、增加 开盘与发行价差距百分比
    df['Gap Open Pct'] = df['Offer_Open_Chng'] / df['Offer_Price']
    df['Gap Open Pct'] = df['Gap Open Pct'].apply(lambda x: 1 if x == np.inf else x)

    # # 7、增加收盘与发行价涨跌幅百分比(未来函数)
    # df['Gap Close Pct'] = df['Offer_Close_Chng'] / df['Offer_Price']

    # print(df.info())
    return df


def explore(df):
    """
    探索数据
    :param df:
    :return:
    """
    # matplotlib中文显示方块
    mpl.rcParams['font.sans-serif'] = ['SimHei']  # 指定默认字体
    mpl.rcParams['axes.unicode_minus'] = False  # 解决保存图像是负号'-'显示为方块的问题

    # df['Year'] = df['Date'].apply(lambda x: x.year)
    # benifit_avg = df.groupby('Year')['Day_Chng'].mean()

    # Day_Chng： （收盘 - 发行价） / 发行价
    # benifit_avg = df.groupby(df['Date'].dt.year)['Day_Chng'].mean()
    # benifit_avg.plot(kind='bar', figsize=(15, 10), color='k', title='IPO首日变化率均值', fontsize=20)

    # benifit_median = df.groupby(df['Date'].dt.year)['Day_Chng'].median()
    # benifit_median.plot(kind='bar', figsize=(15, 10), color='k', title='IPO首日变化率中位数', fontsize=20)

    # print(df['Day_Chng'].describe())
    #
    # df['Day_Chng'].hist(figsize=(15, 7), bins=50, color='grey')

    # Chng： （收盘 - 开盘价） / 开盘价
    df['Chng'] = (df['Close'] - df['Open']) / df['Open']
    print(df['Chng'].describe())
    print(df[df['Chng'] > 0]['Chng'].describe())

    df['Chng'].hist(figsize=(15, 7), bins=50, color='grey')

    plt.show()


def model(df):
    """
    建模
    特征：
    月份：Month
    周：Week
    主要承销商：Lead Mgr
    承销商数量：Total Underwriters
    发行价到开盘价的差距百分比：Gap Open Pct
    发行价到开盘价的价差：Offer_Open_Chng
    发行价：Offer_Price
    开盘价：Open
    sp500前一天从收盘到开盘的变化百分数: Pct
    sp500前一周的变化: Pct_7

    预测结果：发行股是否值得购买（逻辑回归）

    逻辑回归公式：ln(p / (1 - p)) = A*X
    :return:
    """
    # matplotlib中文显示方块
    mpl.rcParams['font.sans-serif'] = ['SimHei']  # 指定默认字体
    mpl.rcParams['axes.unicode_minus'] = False  # 解决保存图像是负号'-'显示为方块的问题

    # 1、数据准备
    df = df.sort_index()
    df['Chng'] = df['Close'] - df['Open']
    print(df['Chng'].describe())
    # print(df.ix[154:155])
    # raise Exception

    # 将各个特征进行稀疏矩阵化
    formula = "Month + Q('Week') + Q('Lead Mgr') + Q('Total Underwriters') + Q('Gap Open Pct') + Q('Offer_Open_Chng') " \
              "+ Q('Offer_Price') + Q('Open') + Q('Pct') + Q('Pct_7')"

    x = dmatrix(formula, data=df, return_type='dataframe')
    df_y = df.loc[x.index]

    # 阈值参数
    y = df_y['Chng'].apply(lambda x: 1 if x > 0.25 else 0)

    x_train, x_test = x.loc[155:], x.loc[:154]
    y_train, y_test = y.loc[155:], y.loc[:154]

    # 2、模型相关
    def logistic():
        """
        逻辑回归
        :return:
        """
        # 1、建模
        clf = linear_model.LogisticRegression()
        clf.fit(x_train, y_train)

        # 2、模型评估
        model_score = clf.score(x_test, y_test)
        print('模型得分：{}'.format(model_score))

        # 3、评估测试集
        pred_label = clf.predict(x_test)
        results = []
        for p1, t1, idx, chg in zip(pred_label, y_test, y_test.index, df_y.loc[y_test.index]['Chng']):
            if p1 == t1:
                results.append([idx, chg, p1, t1, 1])
            else:
                results.append([idx, chg, p1, t1, 0])

        rf = pd.DataFrame(results, columns=['index', 'chng', 'pred', 'actual', 'correct'])
        print(rf[rf['pred'] == 1]['chng'].describe())

        # 5、画图
        fig, ax = plt.subplots(figsize=(15, 10))
        rf[rf['pred'] == 1]['chng'].plot(kind='bar')
        ax.set_xlabel('Index', fontsize=15)
        ax.set_ylabel('当日绝对收益', fontsize=15)
        ax.set_title('预测买入模型', fontsize=20)

        plt.show()

        # 7、由逻辑回归，查看其拟合参数
        fv = pd.DataFrame(clf.coef_).T
        fv.columns = ['Coef']
        fv['Feature'] = pd.DataFrame(x_train.columns, index=fv.index)
        fv = fv.sort_values('Coef', ascending=False).reset_index(drop=True)
        print(fv.head())

    def random_forest():
        """
        随机森林
        对特征重要性排名
        :return:
        """
        # 1、建模
        clf = RandomForestClassifier(n_estimators=1000)
        clf.fit(x_train, y_train)

        # 2、模型评估
        model_score = clf.score(x_test, y_test)
        print('模型得分：{}'.format(model_score))

        feature_importances = clf.feature_importances_
        feature_names = x_train.columns
        feature_std = np.std([tree.feature_importances_ for tree in clf.estimators_], axis=0)

        zz = zip(feature_importances, feature_names, feature_std)
        zzs = sorted(zz, key=lambda x: x[0], reverse=True)

        imps = [x[0] for x in zzs[: 20]]
        labels = [x[1] for x in zzs[: 20]]
        errors = [x[2] for x in zzs[: 20]]

        # 3、画图
        plt.subplots(figsize=(15, 10))
        plt.bar(range(20), imps, color='r', yerr=errors, align='center')
        plt.xticks(range(20), labels, rotation=-70, fontsize=15)

        plt.show()

    # print(df_y['Chng'].loc[:154].describe())
    # 逻辑回归
    # clf = logistic()
    # 随机森林
    clf = random_forest()

    # 6、检测模型鲁棒性，调参
    # 1）将阈值从1美元调整到0.25美元
    # 2) 调整训练集+测试集范围


def run():
    """
    主程序,入口程序
    :return:
    """
    # 1、获取数据
    df_ipo = get_data()

    # print(df_ipo.head(1).T)

    # 2、探索数据
    # explore(df_ipo)

    # 3、建模
    model(df_ipo)


if __name__ == '__main__':
    run()
